{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "66398b75",
   "metadata": {},
   "source": [
    "# Multiple Retrieval Sources\n",
    "\n",
    "Often times you may want to do retrieval over multiple sources. These can be different vectorstores (where one contains information about topic X and the other contains info about topic Y). They could also be completely different databases altogether!\n",
    "\n",
    "A key part is is doing as much of the retrieval in parrelel as possible. This will keep the latency as low as possible. Luckily, [LangChain Expression Language](../../) supports parrellism out of the box.\n",
    "\n",
    "Let's take a look where we do retrieval over a SQL database and a vectorstore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1c5bab6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43a6210f",
   "metadata": {},
   "source": [
    "## Set up SQL query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ab3bf8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.utilities import SQLDatabase\n",
    "from langchain.chains import create_sql_query_chain\n",
    "\n",
    "db = SQLDatabase.from_uri(\"sqlite:///../../../../../notebooks/Chinook.db\")\n",
    "query_chain = create_sql_query_chain(ChatOpenAI(temperature=0), db)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8585120",
   "metadata": {},
   "source": [
    "## Set up vectorstore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b916b0b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.indexes import VectorstoreIndexCreator\n",
    "from langchain.schema.document import Document\n",
    "index_creator = VectorstoreIndexCreator()\n",
    "index = index_creator.from_documents([Document(page_content=\"Foo\")])\n",
    "retriever = index.vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b91816",
   "metadata": {},
   "source": [
    "## Combine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4423211c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import ChatPromptTemplate\n",
    "\n",
    "system_message = \"\"\"Use the information from the below two sources to answer any questions.\n",
    "\n",
    "Source 1: a SQL database about employee data\n",
    "<source1>\n",
    "{source1}\n",
    "</source1>\n",
    "\n",
    "Source 2: a text database of random information\n",
    "<source2>\n",
    "{source2}\n",
    "</source2>\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages([(\"system\", system_message), (\"human\", \"{question}\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7ff87e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_chain = {\n",
    "    \"source1\": {\"question\": lambda x: x[\"question\"]} | query_chain | db.run,\n",
    "    \"source2\": (lambda x: x['question']) | retriever,\n",
    "    \"question\": lambda x: x['question'],\n",
    "} | prompt | ChatOpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d6706410",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 4 is greater than number of elements in index 1, updating n_results = 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "content='There are 8 employees.' additional_kwargs={} example=False\n"
     ]
    }
   ],
   "source": [
    "response = full_chain.invoke({\"question\":\"How many Employees are there\"})\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
